from flask import Flask, request, jsonify
from flask_cors import CORS
from flask_socketio import SocketIO, emit, join_room, leave_room
import pymysql
import random
import smtplib
from email.mime.text import MIMEText
from email.header import Header
import hashlib
import time
import jwt
import joblib
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
from sklearn.preprocessing import MinMaxScaler
import os

app = Flask(__name__)
CORS(app, resources={r"/*": {"origins": "*"}})

SECRET_KEY = os.getenv("SECRET_KEY", "your-secret-key")
app.config['SECRET_KEY'] = SECRET_KEY
socketio = SocketIO(app, cors_allowed_origins="*")

EXPERT_USERNAME = "123456"

def get_db_connection():
    return pymysql.connect(
        host=os.getenv("DB_HOST", "localhost"),
        user=os.getenv("DB_USER", "root"),
        password=os.getenv("DB_PASSWORD", "123456"),
        database=os.getenv("DB_NAME", "dachuang"),
        charset='utf8mb4',
        cursorclass=pymysql.cursors.DictCursor
    )

def generate_verification_code():
    return ''.join([str(random.randint(0, 9)) for _ in range(6)])

def send_verification_email(email, code):
    sender = os.getenv("EMAIL_SENDER", "952304392@qq.com")
    password = os.getenv("EMAIL_PASSWORD", "iklmdpxrhkzpbeha")

    msg = MIMEText(f'您的验证码是：{code}，有效期20分钟', 'plain', 'utf-8')
    msg['From'] = Header(sender)
    msg['To'] = Header(email)
    msg['Subject'] = Header('验证码')

    try:
        smtp_obj = smtplib.SMTP_SSL('smtp.qq.com', 465)
        smtp_obj.login(sender, password)
        smtp_obj.sendmail(sender, [email], msg.as_string())
        smtp_obj.quit()
        return True
    except Exception as e:
        print(f"邮件发送失败: {e}")
        return False

def hash_password(password):
    return hashlib.sha256(password.encode()).hexdigest()

def generate_token(user_id, username):
    try:
        payload = {
            'user_id': user_id,
            'username': username,
            'exp': datetime.utcnow() + timedelta(hours=24)
        }
        token = jwt.encode(payload, SECRET_KEY, algorithm='HS256')
        return token
    except Exception as e:
        print(f"生成 token 失败: {str(e)}")
        raise

def verify_token(token):
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=['HS256'])
        return payload
    except jwt.ExpiredSignatureError:
        return None
    except jwt.InvalidTokenError:
        return None

@app.route('/register/code', methods=['POST'])
def register_code():
    data = request.get_json()
    username = data.get('username')
    email = data.get('email')

    if not all([username, email]):
        return jsonify({'message': '缺少必要参数'}), 400

    conn = get_db_connection()
    try:
        with conn.cursor() as c:
            c.execute('SELECT * FROM users WHERE username = %s OR email = %s', (username, email))
            if c.fetchone():
                return jsonify({'message': '用户名或邮箱已存在'}), 400

            code = generate_verification_code()

            if send_verification_email(email, code):
                c.execute('INSERT INTO verification_codes (email, code, timestamp) VALUES (%s, %s, %s) '
                          'ON DUPLICATE KEY UPDATE code = %s, timestamp = %s',
                          (email, code, int(time.time()), code, int(time.time())))
                conn.commit()
                return jsonify({'message': '验证码已发送至您的邮箱'})
            else:
                return jsonify({'message': '验证码发送失败'}), 500
    finally:
        conn.close()

@app.route('/register', methods=['POST'])
def register():
    data = request.get_json()
    username = data.get('username')
    password = data.get('password')
    email = data.get('email')
    code = data.get('code')

    if not all([username, password, email, code]):
        return jsonify({'message': '缺少必要参数'}), 400

    conn = get_db_connection()
    try:
        with conn.cursor() as c:
            c.execute('SELECT code, timestamp FROM verification_codes WHERE email = %s', (email,))
            result = c.fetchone()

            if not result:
                return jsonify({'message': '验证码无效'}), 400

            stored_code, timestamp = result['code'], result['timestamp']

            if int(time.time()) - timestamp > 1200:
                return jsonify({'message': '验证码已过期'}), 400

            if stored_code != code:
                return jsonify({'message': '验证码错误'}), 400

            hashed_password = hash_password(password)
            c.execute('INSERT INTO users (username, password, email, verified) VALUES (%s, %s, %s, %s)',
                      (username, hashed_password, email, 1))
            conn.commit()
            c.execute('DELETE FROM verification_codes WHERE email = %s', (email,))
            conn.commit()
            return jsonify({'message': '注册成功'})
    finally:
        conn.close()


@app.route('/login/code', methods=['POST'])
def login_code():
    data = request.get_json()

    email = data.get('email')

    if not email:
        return jsonify({'message': '缺少邮箱参数'}), 400

    conn = get_db_connection()
    try:
        with conn.cursor() as c:
            c.execute('SELECT * FROM users WHERE email = %s', (email,))
            if not c.fetchone():
                return jsonify({'message': '邮箱未注册'}), 400

            code = generate_verification_code()

            if send_verification_email(email, code):
                c.execute('INSERT INTO verification_codes (email, code, timestamp) VALUES (%s, %s, %s) '
                          'ON DUPLICATE KEY UPDATE code = %s, timestamp = %s',
                          (email, code, int(time.time()), code, int(time.time())))
                conn.commit()
                return jsonify({'message': '验证码已发送至您的邮箱'})
            else:
                return jsonify({'message': '验证码发送失败'}), 500
    finally:
        conn.close()

@app.route('/login', methods=['POST'])
def login():
    data = request.get_json()
    username = data.get('username')
    password = data.get('password')
    email = data.get('email')
    code = data.get('code')

    if not all([username, password, email, code]):
        return jsonify({'message': '缺少必要参数'}), 400

    conn = get_db_connection()
    try:
        with conn.cursor() as c:
            c.execute('SELECT id, password, email FROM users WHERE username = %s', (username,))
            result = c.fetchone()

            if not result:
                return jsonify({'message': '用户不存在'}), 400

            user_id, stored_password, stored_email = result['id'], result['password'], result['email']

            if stored_email != email:
                return jsonify({'message': '邮箱不匹配'}), 400

            if hash_password(password) != stored_password:
                return jsonify({'message': '密码错误'}), 400

            c.execute('SELECT code, timestamp FROM verification_codes WHERE email = %s', (email,))
            result = c.fetchone()

            if not result:
                return jsonify({'message': '验证码无效'}), 400

            stored_code, timestamp = result['code'], result['timestamp']

            if int(time.time()) - timestamp > 1200:
                return jsonify({'message': '验证码已过期'}), 400

            if stored_code != code:
                return jsonify({'message': '验证码错误'}), 400

            c.execute('DELETE FROM verification_codes WHERE email = %s', (email,))
            conn.commit()

            token = generate_token(user_id, username)
            return jsonify({
                'message': '登录成功',
                'token': token,
                'user_id': user_id,
                'username': username
            })
    finally:
        conn.close()

@app.route('/users', methods=['GET'])
def get_users():
    token = request.headers.get('Authorization')
    if not token:
        return jsonify({'message': '缺少 token'}), 401
    if token.startswith('Bearer '):
        token = token[7:].strip()
    else:
        return jsonify({'message': '无效的 Authorization 头'}), 401

    payload = verify_token(token)
    if not payload:
        return jsonify({'message': '无效或过期的 token'}), 401

    if payload['username'] != EXPERT_USERNAME:
        return jsonify({'message': '仅专家可访问'}), 403

    conn = get_db_connection()
    try:
        with conn.cursor() as c:
            c.execute('SELECT id, username FROM users WHERE username != %s', (EXPERT_USERNAME,))
            users = c.fetchall()
            return jsonify([{'id': u['id'], 'username': u['username']} for u in users])
    finally:
        conn.close()

@app.route('/messages/<int:receiver_id>', methods=['GET'])
def get_messages(receiver_id):
    token = request.headers.get('Authorization')
    if not token:
        return jsonify({'message': '缺少 token'}), 401
    if token.startswith('Bearer '):
        token = token[7:].strip()
    else:
        return jsonify({'message': '无效的 Authorization 头'}), 401

    payload = verify_token(token)
    if not payload:
        return jsonify({'message': '无效或过期的 token'}), 401

    user_id = payload['user_id']
    username = payload['username']

    conn = get_db_connection()
    try:
        with conn.cursor() as c:
            if username != EXPERT_USERNAME:
                c.execute('SELECT id FROM users WHERE username = %s', (EXPERT_USERNAME,))
                expert = c.fetchone()
                if not expert or receiver_id != expert['id']:
                    return jsonify({'message': '用户只能与专家聊天'}), 403

            c.execute('''
                SELECT sender_id, receiver_id, content, timestamp
                FROM messages
                WHERE (sender_id = %s AND receiver_id = %s) OR (sender_id = %s AND receiver_id = %s)
                ORDER BY timestamp ASC
            ''', (user_id, receiver_id, receiver_id, user_id))
            messages = c.fetchall()
            return jsonify([{
                'sender_id': m['sender_id'],
                'receiver_id': m['receiver_id'],
                'content': m['content'],
                'timestamp': m['timestamp'].isoformat()
            } for m in messages])
    finally:
        conn.close()

@app.route('/upload_financial_data', methods=['POST'])
def upload_financial_data():
    token = request.headers.get('Authorization')
    if not token:
        return jsonify({'message': '缺少 token'}), 401
    if token.startswith('Bearer '):
        token = token[7:].strip()
    else:
        return jsonify({'message': '无效的 Authorization 头'}), 401

    payload = verify_token(token)
    if not payload:
        return jsonify({'message': '无效或过期的 token'}), 401

    user_id = payload['user_id']
    data = request.get_json()

    required_fields = [
        'averageAccountsReceivable', 'netCreditSales', 'netProfit',
        'operatingRevenue', 'flowLiabilities', 'totalLiabilities', 'totalAssets'
    ]
    if not all(field in data for field in required_fields):
        return jsonify({'message': '缺少必要参数'}), 400

    try:
        avg_accounts_receivable = float(data['averageAccountsReceivable'])
        net_credit_sales = float(data['netCreditSales'])
        net_profit = float(data['netProfit'])
        operating_revenue = float(data['operatingRevenue'])
        flow_liabilities = float(data['flowLiabilities'])
        total_liabilities = float(data['totalLiabilities'])
        total_assets = float(data['totalAssets'])

        conn = get_db_connection()
        try:
            with conn.cursor() as c:
                c.execute('''
                    INSERT INTO financial_metrics (
                        user_id, average_accounts_receivable, net_credit_sales,
                        net_profit, operating_revenue, flow_liabilities,
                        total_liabilities, total_assets
                    ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
                ''', (
                    user_id, avg_accounts_receivable, net_credit_sales,
                    net_profit, operating_revenue, flow_liabilities,
                    total_liabilities, total_assets
                ))
                conn.commit()
                return jsonify({'message': '财务数据上传成功'})
        finally:
            conn.close()

    except Exception as e:
        return jsonify({'message': f'上传失败：{str(e)}'}), 500

@app.route('/calculate_financial_metrics', methods=['POST'])
def calculate_financial_metrics():
    token = request.headers.get('Authorization')
    if not token:
        return jsonify({'message': '缺少 token'}), 401
    if token.startswith('Bearer '):
        token = token[7:].strip()
    else:
        return jsonify({'message': '无效的 Authorization 头'}), 401

    payload = verify_token(token)
    if not payload:
        return jsonify({'message': '无效或过期的 token'}), 401

    user_id = payload['user_id']

    try:
        conn = get_db_connection()
        try:
            with conn.cursor() as c:
                c.execute('''
                    SELECT timestamp, average_accounts_receivable, net_credit_sales,
                           net_profit, operating_revenue, flow_liabilities,
                           total_liabilities, total_assets
                    FROM financial_metrics
                    WHERE user_id = %s
                    ORDER BY timestamp DESC
                    LIMIT 5
                ''', (user_id,))
                history = c.fetchall()

                if not history:
                    return jsonify({'message': '没有历史数据可用于生成图表'}), 404

                labels = [row['timestamp'].strftime('%Y-%m-%d %H:%M') for row in history]
                accounts_receivable_turnover = []
                net_profit_margin = []
                debt_ratio_flow = []
                debt_ratio_nonflow = []
                asset_liability_ratio = []
                operating_revenue_data = []

                for row in history:
                    avg_accounts_receivable = float(row['average_accounts_receivable'] or 0)
                    net_credit_sales = float(row['net_credit_sales'] or 0)
                    net_profit = float(row['net_profit'] or 0)
                    operating_revenue = float(row['operating_revenue'] or 0)
                    flow_liabilities = float(row['flow_liabilities'] or 0)
                    total_liabilities = float(row['total_liabilities'] or 0)
                    total_assets = float(row['total_assets'] or 0)

                    accounts_receivable_turnover.append(
                        (net_credit_sales / avg_accounts_receivable) if avg_accounts_receivable != 0 else 0
                    )
                    net_profit_margin.append(
                        (net_profit / operating_revenue * 100) if operating_revenue != 0 else 0
                    )
                    non_flow_liabilities = total_liabilities - flow_liabilities
                    debt_ratio_flow.append(
                        (flow_liabilities / total_liabilities * 100) if total_liabilities != 0 else 0
                    )
                    debt_ratio_nonflow.append(
                        (non_flow_liabilities / total_liabilities * 100) if total_liabilities != 0 else 0
                    )
                    asset_liability_ratio.append(
                        (total_liabilities / total_assets * 100) if total_assets != 0 else 0
                    )
                    operating_revenue_data.append(operating_revenue)

                echarts_data = {
                    'echarts1': {
                        'xAxis': labels,
                        'series': [{
                            'name': '应收账款周转率',
                            'type': 'line',
                            'data': [round(x, 2) for x in accounts_receivable_turnover]
                        }]
                    },
                    'echarts2': {
                        'xAxis': labels,
                        'series': [{
                            'name': '营业收入',
                            'type': 'line',
                            'data': [round(x, 2) for x in operating_revenue_data]
                        }]
                    },
                    'echarts3': {
                        'xAxis': labels,
                        'series': [{
                            'name': '净利润率',
                            'type': 'bar',
                            'data': [round(x, 2) for x in net_profit_margin]
                        }]
                    },
                    'echarts4': {
                        'series': [{
                            'name': '负债比',
                            'type': 'pie',
                            'data': [
                                {'value': round(debt_ratio_flow[0], 2), 'name': '流动负债'},
                                {'value': round(debt_ratio_nonflow[0], 2), 'name': '非流动负债'}
                            ]
                        }]
                    },
                    'echarts5': {
                        'yAxis': labels,
                        'series': [{
                            'name': '资产负债率',
                            'type': 'bar',
                            'data': [round(x, 2) for x in asset_liability_ratio]
                        }]
                    }
                }

                return jsonify(echarts_data)

        finally:
            conn.close()

    except Exception as e:
        return jsonify({'message': f'计算失败：{str(e)}'}), 500

@app.route('/delete_financial_data', methods=['POST'])
def delete_financial_data():
    token = request.headers.get('Authorization')
    if not token:
        return jsonify({'message': '缺少 token'}), 401
    if token.startswith('Bearer '):
        token = token[7:].strip()
    else:
        return jsonify({'message': '无效的 Authorization 头'}), 401

    payload = verify_token(token)
    if not payload:
        return jsonify({'message': '无效或过期的 token'}), 401

    user_id = payload['user_id']

    conn = get_db_connection()
    try:
        with conn.cursor() as c:
            c.execute('DELETE FROM financial_metrics WHERE user_id = %s', (user_id,))
            conn.commit()

            if c.rowcount > 0:
                return jsonify({'message': '财务数据删除成功'})
            else:
                return jsonify({'message': '没有找到可删除的财务数据'})
    except Exception as e:
        return jsonify({'message': f'删除失败：{str(e)}'}), 500
    finally:
        conn.close()

@app.route('/predict_risk', methods=['POST'])
def predict_risk():
    data = request.get_json()
    required_fields = [
        'debtToAssetRatio', 'netProfitMargin', 'cashFlowCoverageRatio',
        'receivablesToRevenueRatio', 'operatingCashFlowRevenueRatio'
    ]

    if not all(field in data for field in required_fields):
        return jsonify({'message': '缺少必要参数'}), 400

    input_data = {k: data[k] for k in required_fields}

    try:
        model_dict = joblib.load('风控模型_v1.pkl')
        field_mapping = {
            'debtToAssetRatio': '资产负债率',
            'netProfitMargin': '净利润率',
            'cashFlowCoverageRatio': '现金流量比率',
            'receivablesToRevenueRatio': '应收账款周转率(次)',
            'operatingCashFlowRevenueRatio': '经营净现金流/营业总收入'
        }
        df = pd.DataFrame([input_data]).rename(columns=field_mapping)
        for col in model_dict['percent_cols']:
            df[col] = df[col] / 100
        y_proba = model_dict['model'].predict_proba(df)[:, 1][0]
        y_pred = 1 if y_proba > 0.5 else 0
        scaler = MinMaxScaler(feature_range=(0, 100))
        risk_index = scaler.fit_transform([[y_proba]])[0][0]
        warning = "红色预警" if risk_index < 60 else "黄色预警" if risk_index < 70 else "正常"

        return jsonify({
            '风险概率': round(y_proba, 4),
            '风险标签': '风险' if y_pred == 1 else '安全',
            '风险指数': round(risk_index, 2),
            '预警信息': warning
        })

    except Exception as e:
        return jsonify({'message': f'预测失败：{str(e)}'}), 500

@socketio.on('connect')
def handle_connect():
    token = request.args.get('token')
    if not token:
        return False
    payload = verify_token(token)
    if not payload:
        return False
    emit('status', {'message': '连接成功'})

@socketio.on('join')
def on_join(data):
    token = data.get('token')
    if not token:
        emit('error', {'message': '缺少 token'})
        return
    payload = verify_token(token)
    if not payload:
        emit('error', {'message': '无效或过期的 token'})
        return
    user_id = payload['user_id']
    room = f'user_{user_id}'
    join_room(room)
    emit('status', {'message': f'用户 {user_id} 已加入房间 {room}'}, room=room)

@socketio.on('leave')
def on_leave(data):
    token = data.get('token')
    if not token:
        emit('error', {'message': '缺少 token'})
        return
    payload = verify_token(token)
    if not payload:
        emit('error', {'message': '无效或过期的 token'})
        return
    user_id = payload['user_id']
    room = f'user_{user_id}'
    leave_room(room)
    emit('status', {'message': f'用户 {user_id} 已离开房间 {room}'}, room=room)

@socketio.on('message')
def handle_message(data):
    token = data.get('token')
    if not token:
        emit('error', {'message': '缺少 token'})
        return
    payload = verify_token(token)
    if not payload:
        emit('error', {'message': '无效或过期的 token'})
        return

    sender_id = payload['user_id']
    username = payload['username']
    receiver_id = data.get('receiver_id')
    content = data.get('content')

    if not all([receiver_id, content]):
        emit('error', {'message': '缺少必要参数'})
        return

    conn = get_db_connection()
    try:
        with conn.cursor() as c:
            c.execute('SELECT id FROM users WHERE id = %s', (receiver_id,))
            if not c.fetchone():
                emit('error', {'message': '接收者不存在'})
                return

            if username != EXPERT_USERNAME:
                c.execute('SELECT id FROM users WHERE username = %s', (EXPERT_USERNAME,))
                expert = c.fetchone()
                if not expert or receiver_id != expert['id']:
                    emit('error', {'message': '用户只能与专家聊天'})
                    return

            c.execute('INSERT INTO messages (sender_id, receiver_id, content) VALUES (%s, %s, %s)',
                      (sender_id, receiver_id, content))
            conn.commit()

            c.execute('SELECT timestamp FROM messages WHERE id = LAST_INSERT_ID()')
            timestamp = c.fetchone()['timestamp']

            receiver_room = f'user_{receiver_id}'
            sender_room = f'user_{sender_id}'
            message_data = {
                'sender_id': sender_id,
                'receiver_id': receiver_id,
                'content': content,
                'timestamp': timestamp.isoformat()
            }
            emit('message', message_data, room=receiver_room)
            emit('message', message_data, room=sender_room)
    finally:
        conn.close()

if __name__ == '__main__':
    socketio.run(app, debug=True, host='0.0.0.0', port=5000, allow_unsafe_werkzeug=True)